#!/usr/bin/env python3

# backup

import os
import sys
import glob
import time
import shutil
from subprocess import Popen, PIPE

err_code_msg = {
        'err_conf_notfond': (1, 'vport conf file not fond.'),
        'err_conf_syntax': (2, 'vport conf file syntax error.'),
        'err_del_vport': (3, 'vport abnormal deletion.'),
        'err_after_reload': (4, 'lvsconf different from ipvsadm after reload.'),
        'err_before_check': (5, 'lvsconf different from ipvsadm before check.'),
        'err_all_conf_syntax': (6, 'syntax error for agent reload.'),
        'err_all_del_vport': (7, 'vport abnormal deletion for agent reload.')
}


sorry_server = {}

def remove_files(fromdir):
        files = glob.glob(os.path.join(fromdir, '*'))
        for filename in files:
                os.remove(filename)


def copy_files(fromdir, todir):
        files = glob.glob(os.path.join(fromdir, '*.cfg'))
        for filename in files:
                shutil.copy(filename, todir)


def goto_err(err_type):
        err_info = err_code_msg[err_type]
        print('Error: %s' % err_info[1])
        sys.exit(err_info[0])


def get_rs_info(filename):
        rs_list = []
        fd = open(filename, 'r')
        for line in fd:
                if line.strip().startswith('#'): continue
                if 'real_server' not in line.strip(): continue
                rip, rport = line.split()[1:3]
                rs_list.append(rip + ':' + rport)
        fd.close()
        rs_list.sort()
        return rs_list


def get_sorry_info(filename):
        sorry = ''
        fd = open(filename, 'r')
        for line in fd:
                if line.strip().startswith('#'): continue
                if 'sorry_server' not in line.strip(): continue
                sip, sport = line.strip().split()[1:3]
                sorry = sip + ':' + sport
        fd.close()
        return sorry
        

def get_lvs_vip_info(dirname):
        ret = {}
        for path in glob.glob(os.path.join(dirname, '*-*.cfg')):
                filename = os.path.basename(path)
                vip      = filename.split('-')[0]
                vport    = filename.split('-')[-1].strip('.cfg')
                rs       = get_rs_info(os.path.join(dirname, vip + '-' + vport + '.cfg'))
                if len(rs) > 0: ret[(vip, vport)] = rs

                # sorry_server info
                sorry = get_sorry_info(os.path.join(dirname, vip + '-' + vport + '.cfg'))
                if sorry != '': sorry_server[(vip, vport)] = sorry
        return ret


def get_ssl_vip_info(dirname):
        # 1. first get ssl_rs info
        fdrs = open('ssl_rs.conf', 'r')
        sslrs_list = []
        for line in fdrs:
                if line.strip().startswith('#') or 'real_server' not in line: continue
                rip, rport = line.strip().split()[1:3]
                sslrs_list.append(rip + ':' +  rport)
        fdrs.close()
        sslrs_list.sort()

        # 2. get ssl vip-vport info
        ret = {}
        for path in glob.glob(os.path.join(dirname, '*-*.cfg')):
                filename = os.path.basename(path)
                vip      = filename.split('-')[0]
                vport    = filename.split('-')[-1].strip('.cfg')
                ret[(vip, vport)] = sslrs_list
        return ret


def get_ipvsadm_vip_info():
        ret = {}
        p = Popen(['/sbin/ipvsadm', '-ln'], stdout = PIPE)
        data = p.stdout.read().decode()
        tcp_list = data.split('TCP')[1:]
        for vs in tcp_list:
                line_list = vs.split('\n')
                vip = vport = ''
                rs_list = []
                for line in line_list:
                        if line == '': continue
                        if not line.strip().startswith('->'):
                                vip, vport = line.strip().split()[0].split(':')
                        else:
                                rs_list.append(line.strip().split()[1])
                rs_list.sort()
                ret[(vip, vport)] = rs_list
        return ret


def compare_lvsconf_ipvsadm_rs():
        def is_list_in(a, b):
                ret = [False for i in a if i not in b]
                return True if len(ret) == 0 else False 
        lvs_dir = 'lvs_file'
        ssl_dir = 'ssl_file'
        lvs_dict = get_lvs_vip_info(lvs_dir)
        ssl_dict = get_ssl_vip_info(ssl_dir)
        lvs_dict.update(ssl_dict)
        
        ipvsadm = get_ipvsadm_vip_info()
        if len(lvs_dict) != len(ipvsadm): return False
        for k,v in ipvsadm.items():
                if ipvsadm[k] != lvs_dict[k]:
                        diff = [i for i in ipvsadm[k] if i not in lvs_dict[k]]
                        if not (len(ipvsadm[k]) == (len(lvs_dict[k]) + 1) and is_list_in(lvs_dict[k], ipvsadm[k]) and diff[0] == sorry_server[k]):
                                return False
        return True

def _slb_api():
        return True


def check_agent_restart():
        group_dir = 'group'
        for path in glob.glob(os.path.join(group_dir, '*-*.cfg')):
                #filename = os.path.basename(path)
                fd = open(path, 'r')
                data = fd.read()
                fd.close()
                if 'real_server' not in data and not _slb_api():
                        return False
        return True


def check_before_reload(agent_param):
        if not compare_lvsconf_ipvsadm_rs():
                goto_err('err_before_check')

        if agent_param == 'all':
                # 1. syntax check
                ret = os.system('./keepalived -tf "group/*-*.cfg" &>/dev/null')
                if ret != 0:
                        goto_err('err_all_conf_syntax')

                # 2. check delete
                if check_agent_restart():
                        # 1.1 delete lvs_file/*
                        ##### remove_files('lvs_file')
                        # 1.2 copy group/* lvs_file/
                        ##### copy_files('group', 'lvs_file')
                        print('copy agent dir to lvs_file...')
                else:
                        goto_err('err_all_del_vport')
        elif '/' in agent_param:
                # 1. syntax check
                groupid, vip, vport = agent_param.split('/')
                filename = os.path.join('group', vip + '-' + vport + '.cfg')
                if not os.path.exists(filename): goto_err('err_conf_notfond')
                ret = os.system('./keepalived -tf %s &>/dev/null' % filename)
                if ret != 0:
                        goto_err("err_conf_syntax")
                # 2. exception ?
                fd = open(filename, 'r')
                data = fd.read()
                if 'real_server' not in data and not _slb_api():
                        goto_err('err_del_vport')
                # 3. normal
                # /bin/cp group/vip-vport.conf lvs_file/
                #### shutil.copy(filename, 'lvs_file/')
                print('copy file %s to lvs_file dir.' % filename)
                fd.close()
                pass
        else:
                print('param error.')
                sys.exit(100)
                pass


def check_after_reload():
        if not compare_lvsconf_ipvsadm_rs():
                goto_err('err_after_reload')
        print('Reload successfully!')
        sys.exit(0)


if __name__ == '__main__':
        # change cwd
        os.chdir('/etc/keepalived')

        if len(sys.argv) == 2 and sys.argv[1] == 'after':
                check_after_reload()
        elif len(sys.argv) == 3 and sys.argv[1] == 'before':
                check_before_reload(sys.argv[2])
        else:
                print('param error.')
                sys.exit(100)
        # kill -HUP pid  :  debug
        time.sleep(1)
        check_after_reload()

